"""
渲染器
"""
import numpy as np
import taichi as ti
# pylint: disable=wildcard-import,unused-wildcard-import,consider-using-f-string

from .gtaichi import (Ray,circle_random_samples)
from .scenes import FilmCamera
from .integrators import RayTraceIntegrator

@ti.data_oriented
class Renderer:
    """
    渲染器框架
    """
    def __init__(self,w=None,h=None,bounding=None,
                 integrator=RayTraceIntegrator):
        # pylint: disable=no-member
        self.bounding = bounding
        self.integrator = integrator(self)
        self.resolution = (800 if w is None else w,450 if h is None else h)
        self.pixels = ti.Vector.field(3, float, self.resolution) # 图像的像素场
        self.hdr_pxiels = ti.Vector.field(3, float, self.resolution)
        self.samples_count = ti.field(dtype=ti.f32,shape=())
        self.aspect_ratio = self.resolution[0]/self.resolution[1]
        self.fov = np.pi/12
        self.fov_tan = ti.field(dtype=ti.f32,shape=())
        self.fov_tan[None] = ti.math.tan(self.fov/2)
        self.hdr = False
        self.lens_l = ti.field(dtype=ti.f32,shape=())
        self.lens_l[None] = 0.
        self.lens_phi = ti.field(dtype=ti.f32,shape=())
        self.lens_phi[None] = 0.
        self.cam_mat4  = ti.math.mat4.field(shape=())

    @ti.kernel
    def fill_screen(self,c:ti.math.vec3):
        """
        填充屏幕颜色
        """
        for i,j in self.pixels:
            self.pixels[i,j] = c

    def clear(self):
        """
        重置缓冲器
        """
        self.samples_count[None] = 0
        self.fill_screen(ti.Vector([0.,0.,0.]))

    def reset_framebuffer(self):
        """
        重新渲染
        """
        self.samples_count[None] = 0
        self.fill_screen(ti.Vector([0.,0.,0.]))

    def accumulate(self):
        """
        累加渲染
        """
        self.render()

    @ti.kernel
    def hdr_pixels(self):
        """
        HDR处理
        """
        for i,j in self.pixels:
            c = self.pixels[i,j]
            self.hdr_pxiels[i,j] = ti.math.sqrt(c / (1. + c))

    def fetch_image(self):
        """
        返回渲染images
        """
        if self.hdr:
            self.hdr_pixels()
            return self.hdr_pxiels
        return self.pixels

    def set_camera(self,cam):
        """
        矩阵mat从世界向摄像机坐标系变换
        """
        if isinstance(cam,FilmCamera):
            self.lens_l[None],self.lens_phi[None] = cam.get_lens()
        self.cam_mat4[None] = ti.Matrix(np.linalg.inv(cam.mat4))

    @ti.func
    def eye_ray(self,u,v):
        """
        返回从眼睛到u,v的射线
        """
        o = ti.Vector([0.,0.,0.,1.])
        p = ti.Vector([0.,0.,0.,1.])
        if self.lens_l[None]==0.:
            z = -0.5/self.fov_tan[None]
            crs = circle_random_samples(1)
            uu = u + crs.x - 0.5
            vv = v + crs.y - 0.5
            p = ti.Vector([self.aspect_ratio*(uu/self.resolution[0]-0.5),
                                            vv/self.resolution[1]-0.5,
                                            z,1.])
        else:
            crs = circle_random_samples(self.lens_phi[None])
            o = ti.Vector([crs.x,crs.y,0.,1.])
            h = 2. * self.lens_l[None] * self.fov_tan[None]
            w = self.aspect_ratio * h
            p = ti.Vector([w/self.resolution[0] * u - w/2.,
                           h/self.resolution[1] * v - h/2.,
                           -self.lens_l[None],1.])
        o = self.cam_mat4[None] @ o
        p = self.cam_mat4[None] @ p
        return Ray(o.xyz,ti.math.normalize(p-o).xyz)

    @ti.kernel
    def render(self):
        """
        渲染场景
        """
        self.samples_count[None] += 1
        s = 1./self.samples_count[None]
        for u,v in self.pixels:
            ray = self.eye_ray(u,v)
            c = self.integrator.l_i(ray)
            self.pixels[u,v] = self.pixels[u,v] * (1.-s) + c * s
